package com.algorithm.Tree;

import printer.BinaryTreeInfo;
import printer.BinaryTrees;

import java.util.Comparator;

public class BinarySearchTree<E> extends BinaryTree<E>{
    private Comparator<E> comparator;

    public BinarySearchTree() {
        this(null);
    }

    public BinarySearchTree(Comparator<E> comparator) {
        this.comparator = comparator;
    }

    public boolean contains(E element) {
        return node(element) != null;
    }

    public void clear() {
        root = null;
        size = 0;
    }

    /**
     * 添加一个元素
     */
    public void add(E element) {
        elementNotNullCheck(element);
        if (root == null) {
            root = createNode(element,null);
            size++;
            afterAdd(root);
            return;
        }

        Node<E> node = root;
        Node<E> parent;
        int cmp;
        do {
            cmp = compare(element,node.element);
            parent = node;
            if (cmp > 0) {
                node = node.right;
            }else if (cmp < 0) {
                node = node.left;
            }else {
                node.element = element;
                return;
            }
        }while (node != null);

        Node<E> newNode = createNode(element,parent);
        if (cmp > 0 ) {
            parent.right = newNode;
        }else {
            parent.left = newNode;
        }
        afterAdd(newNode);
        size++;
    }

    /**
     * remove element
     */
    public void remove(E element) {
        if (root == null) return;
        Node<E> node = node(element);
        if (node == null) return;

        if (node.hasTwoChild()) {
            Node<E> pre = successor(node);
            node.element = pre.element;
            node = pre;
        }

        Node<E> replaceNode = node.left != null ? node.left : node.right;
        if (replaceNode != null) {
            replaceNode.parent = node.parent;
            if (node.parent == null) {
                root = replaceNode;
            }else if (node.isLeftChild()) {
                node.parent.left = replaceNode;
            }else {
                node.parent.right = replaceNode;
            }
            afterRemove(replaceNode);
        }else if (node.parent == null) {
            root = null;
        }else {
            if (node.isLeftChild()) {
                node.parent.left = null;
            }else {
                node.parent.right = null;
            }
            afterRemove(node);
        }
        size--;
    }

    protected void afterAdd(Node<E> node) {}

    protected void afterRemove(Node<E> node) {}


    protected Node<E> createNode(E element,Node<E> parent) {
        return new Node<>(element,parent);
    }

    protected int compare(E ele1,E ele2) {
        if (comparator != null) {
            return comparator.compare(ele1,ele2);
        }
        return ((Comparable) ele1).compareTo(ele2);
    }

    protected Node<E> node(E element) {
        if (root == null) return null;
        elementNotNullCheck(element);
        Node<E> node = root;
        while (node != null) {
            int cmp = compare(element,node.element);
            if (cmp == 0) return node;
            if (cmp > 0) {
                node = node.right;
            }else {
                node = node.left;
            }
        }
        return null;
    }

    protected void rotatedLeft(Node<E> grand) {
        Node<E> parent = grand.right;
        Node<E> child = parent.left;

        grand.right = child;
        parent.left = grand;
        afterRotated(grand,parent,child);
    }

    protected void rotatedRight(Node<E> grand) {
        Node<E> parent = grand.left;
        Node<E> child = parent.right;

        grand.left = child;
        parent.right = grand;
        afterRotated(grand,parent,child);
    }

    protected void afterRotated(Node<E> grand,Node<E> parent,Node<E> child) {
        parent.parent = grand.parent;
        if (grand.isLeftChild()) {
            parent.parent.left = parent;
        }else if (grand.isRightChild()) {
            parent.parent.right = parent;
        }else {
            root = parent;
        }

        if (child != null) {
            child.parent = grand;
        }
        grand.parent = parent;
    }


    public void print() {
        BinaryTreeInfo binaryTreeInfo = new BinaryTreeInfo() {
            @Override
            public Object root() {
                return root;
            }

            @Override
            public Object left(Object node) {
                return ((Node) node).left;
            }

            @Override
            public Object right(Object node) {
                return ((Node) node).right;
            }

            @Override
            public Object string(Object node) {
                return node;
            }
        };
        BinaryTrees.println(binaryTreeInfo);
    }
}
